open Printf
open Ext

(* Globals used for command line parameters *)
let mfile = ref ""
let depnet = ref false
let datafile = ref ""
let pll = ref false
let clib = ref true

let usage = "Usage: mscore -m <model> -i <data>"
let args = Arg.align
 ([("-m", Arg.Set_string mfile, " Input model (AC, BIF, XMOD, or MN format)") ;
   ("-depnet", Arg.Set depnet, " Score model as a DN, using pseudo-log-likelihood");
   ("-i", Arg.Set_string datafile, " Test data file");
   ("-pll", Arg.Set pll, " Compute pseudo-log-likelihood");
   ("-clib", Arg.Set clib, " Use C library for computing PLL");
   ("-noclib", Arg.Clear clib, " Use only OCaml code for computing PLL")]
   @ common_arguments)

let main () = 
  Arg.parse args ignore usage;
  if !mfile = "" || !datafile = "" then
    (Arg.usage args usage; exit 0);

  common_log_init();

  let score_one =
    match filetype !mfile with
    | ACFile ->
      (* TODO -- eventually support PLL for circuits, by using the
         loaded MN representation. *)
      if !pll then begin
        nlogf "WARNING: Pseudo-likelihood not supported for circuits.\n";
        nlogf "Using regular log-likelihood instead.\n";
        pll := false
      end;
      let circ = Circuit.load (open_in !mfile) in
      let scratch = Circuit.create_scratch circ in
      let z = Circuit.compute_z scratch circ in
      (fun x -> Circuit.loglikelihood_x scratch circ x -. z)

    | BNFile ->
      let bn = Bn.load_auto !mfile in
      if !pll && not !depnet then Bn.pll bn
      else Bn.loglikelihood bn
    | MNFile ->
      let mn = Mn.load_auto !mfile in
      if !pll && not !depnet then 
        let fa = Array.of_list (Mn.to_features mn) in
        if !clib then
          let cmn = Pll.create_mn mn.Mn.schema fa in
          Pll.pll_mn cmn 
        else
          Pll.fa_pll mn.Mn.schema fa 
      else Mn.raw_logprob mn
    | _ -> printf "ERROR: Could not determine filetype from extension.\n";
           exit (-1) in
  let datastream = open_in !datafile in 
  let s = stats_make () in
  try 
    while true do
      let (w, x) = Data.input_wexample datastream in
      let ll = score_one x in 
      vlogf "%f\n" ll ; 
      stats_wadd s w ll
    done
  with Data.Eof -> 
    nlogf "avg = %f +/- %f\n" (stats_mean s) (stats_stderr s)

;;
main ()
